feat: [WIP/Proposal] Support copick global labels as the persisted label space#74
feat: [WIP/Proposal] Support copick global labels as the persisted label space#74uermel wants to merge 2 commits into
Conversation
daniel-ji
left a comment
There was a problem hiding this comment.
A Claude report, reviewed and edited by me
Verdict: ✅ Sound. Core mechanism verified end-to-end; default path unchanged. Five minor, non-blocking nits.
This report can be found at [czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr/REPORT.md.
Replace [czii-dir] with the actual path to the czii dir on the HPC.
1. Environment
octopi conda env: ~/miniforge3/envs/octopi
pip install -e [czii-dir]/cryoet-data-portal-pick-extract/octopi --no-deps(editable branch install; regenerates theoctopiconsole script).
Test fixture: [czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr/sample_copick_project — a copy of the 31 MB sample_copick_project unit fixture with config roots rewritten to the copy. 3 runs (TS_001–003), 64³ tomograms. Objects: proteasome=1, ribosome=2, membrane=3. Targeting ribosome alone forces divergence (global 2 ≠ dense channel 1). No original copick/octopi files were modified.
2. Test results (all pass)
Tier 0 — 19 unit tests (.../octopi-74-pr/test_tier0_units.py, no GPU): finalize_label_space (mint, divergence, --from-model inherit, missing-object error), build_compaction_map (legacy (None,None); copick orig=[0,1,2,3]/target=[0,0,1,2]), _build_dense_to_global (expansion [0,2,3], legacy identity, and the 4 raises: count mismatch, out-of-range, unknown object, ensemble disagreement). 19/19 passed.
Tier 1 — create-targets on-disk (no GPU), ribosome alone:
| Mode | On-disk seg (TS_001) | YAML labels |
YAML label_space |
|---|---|---|---|
| copick | [0, 2] (ribosome global) |
{ribosome: 1} |
copick |
| model (default) | [0, 1] (dense channel) |
{ribosome: 1} |
(absent) |
get_class_info + build_compaction_map on the copick target: label_space=copick, orig=[0,1,2,3] → target=[0,0,1,0] (disk 2 → channel 1). ✅
Tier 2 — GPU round-trip (A40): train → segment → localize, copick-mode ribosome target, 2 epochs, -dim 64:
model_config.yaml:label_space: copick,labels: {ribosome:1},num_classes: 2. ✅segmentwrotepredict-copick=[0, 2]on all 3 runs (global label, not channel 1) → dense→global expansion confirmed. ✅localizeused objects[['ribosome', 2, 150.0]](global label kept, no legacy remap) and extracted a ribosome pick. ✅
3. Findings (minor, non-blocking)
1. --from-model on an object subset fails late, not fast.
num_classes is derived from the count of selected objects (get_class_info: len(labels)+1), but --from-model keeps each object's inherited channel. Fine-tuning on a subset can leave an inherited channel ≥ the shrunk num_classes.
Example: pretrained {ribosome→1, membrane→2} (3 classes); fine-tune on membrane only → membrane→2 but num_classes=2 (channels {0,1}). Target value 2 is out of range; surfaces at training (one-hot out-of-bounds) or inference ("Model label 2 for 'membrane' is out of range [1, 1]"). Never silently corrupts.
| Path | membrane channel | num_classes |
Result |
|---|---|---|---|
| Default mint, membrane only | 1 | 2 | ✅ |
--from-model, membrane only |
2 | 2 | ❌ value 2 in 2-channel model |
--from-model, membrane + ribosome (empty source) |
2 | 3 | ✅ ribosome = dead channel 1 |
Workaround (no code change): list the full pretrained object set, giving data-less objects an empty pick source (e.g. -t "ribosome,nobody,999" -t "membrane,membrain,0"). They stay in the labels schema, so num_classes is preserved and channels stay contiguous. Verified: labels={ribosome:1, membrane:2}, num_classes=3, compaction [0,1,2,3]→[0,0,1,2], dense_to_global=[0,2,3].
Suggested fix: when --from-model is set, inherit the pretrained config's entire labels map (all objects/channels), so num_classes and channel bindings always match the pretrained model. Alternatively, add a contiguity check in finalize_label_space to fail fast.
2. uint8 ceiling. Training compaction (get_transforms(label_dtype=np.uint8)) and inference expansion (dense_to_global, np.uint8) are uint8; a global label > 255 wraps silently. Unguarded (low risk — copick segs are uint8).
3. MultiCopickDataModule missing None guard (generators.py:388): ref = {... self.roots[first_session].get_object(name).label ...} has no None check, unlike the loop at :391. First-session-missing-object → AttributeError instead of the intended ValueError. Multi-config training only.
4. get_class_info UnboundLocalError on no-seg (helpers.py:270–301): if no target segmentation is found for any runID, the return variables are never assigned. Pre-existing (loop-and-break predates the PR); warrants a clean error.
5. Compaction assumes on-disk seg holds only current pickable globals. MapLabelValued passes through any value not in orig_labels (built from root.pickable_objects). A stray on-disk label (config changed after painting) would enter dense space unmapped and could exceed num_classes. Low probability.
4. Reproduction
source ~/miniforge3/etc/profile.d/conda.sh && conda activate octopi # however you activate your octopi env
BASE=[czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr
CFG=$BASE/sample_copick_project/filesystem.json
# Tier 0
python $BASE/test_tier0_units.py
# Tier 1
octopi create-targets -c $CFG -t "ribosome,gapstop,0" -alg wbp -vs 10 \
-name ribo-copick -uid octopi -sid 1 --label-space copick
octopi create-targets -c $CFG -t "ribosome,gapstop,0" -alg wbp -vs 10 \
-name ribo-model -uid octopi -sid 1
# Tier 2 (GPU)
octopi train -c $CFG --tomo-uri wbp@10.0 --target-uri "ribo-copick:octopi/1" \
--trainRunIDs TS_001,TS_002 --validateRunIDs TS_003 -nepochs 2 -dim 64 -bs 2 \
-o $BASE/model_copick
octopi segment -c $CFG --tomo-uri wbp@10.0 \
--model-config $BASE/model_copick/model_config.yaml \
--model-weights $BASE/model_copick/best_model.pth \
--seg-uri "predict-copick:octopi/1" --run-ids TS_001,TS_002,TS_003
octopi localize -c $CFG --seg-uri "predict-copick:octopi/1" -vs 10 \
--runIDs TS_001,TS_002,TS_003 -pui octopi -psid loc1Verify on-disk labels:
import copick, numpy as np
r = copick.from_file(CFG)
segs = r.get_run("TS_001").get_segmentations(name="predict-copick", user_id="octopi", session_id="1", is_multilabel=True)
print(sorted(np.unique(segs[0].numpy()).tolist())) # -> [0, 2] (global, not channel 1)5. Artifacts (under [czii-dir]/cryoet-data-portal-pick-extract/octopi-74-pr/)
sample_copick_project/— isolated fixture copy (roots rewritten to this dir)test_tier0_units.py— 19 unit testsmodel_copick/— trained model (best_model.pth,model_config.yaml,results.csv)train.log,seg.log,loc.log— run logspre_ribo_mb.yaml,pretrained_model_config.yaml,pretrained_missing.yaml,pre_contig.yaml—--from-modelfixtures
Summary
Octopi has always persisted a model-local, dense, contiguous label space
(
1, 2, 3, …assigned in argument order bycreate-targets) into the copicksegmentations it writes (
targets,predict). This PR adds an opt-inalternative in which the segmentations octopi writes hold copick's global
object labels, while the dense
[0..K]space that the model needs exists onlytransiently inside the training/inference boundary.
The behavior is selected with a new
--label-spaceflag onoctopi create-targets:--label-space model(default) — unchanged behavior.--label-space copick— paint canonical copick global labels on disk; carry a{name: model_channel}map plus alabel_space: copickmarker through thetargets YAML → model config → segment YAML so every stage can convert between
the two spaces by object name.
Rationale
In copick, an integer label means the same object everywhere in a project (e.g.
gas-vesicle-shell = 2). Octopi's sequential relabeling breaks that invariant:a
targets/predictsegmentation stored by octopi might relabelgas-vesicle-shellto1, so the file can only be interpreted if you also carryoctopi's
labelsmap alongside it. Opening such a segmentation with the plainproject config (napari, copick logical/process ops, any downstream consumer)
mislabels every voxel.
--label-space copickfixes this at the source: what octopi writes to disk isordinary, canonical copick multilabel. The model still trains and predicts in a
dense
[0..K]space (required by the loss, the one-hot encoding, andRandCropByLabelClassesd), but that space is now a private, derived detail —compacted from global labels right after loading, and expanded back to global
labels right before writing.
A second benefit: because both conversions are resolved by object name (not
by sort order or by hard-coded integers), a model trained on one project can be
applied to another project whose global labels differ, and the written
predictsegmentation will use the target project's globals.How it works
Two label concepts exist in
copickmode:[0..K]index the network sees, recorded as thelabels: {name: channel}map in the YAMLs.The pipeline converts between them at exactly two boundaries:
build_compaction_map+ MONAIMapLabelValued, inserted rightafter
LoadCopickd): global → dense. The map is derived from the recorded{name: channel}joined with the project's global labels by name. Objects thatare painted but not selected as targets fold to background
0._build_dense_to_global, applied just before the single copickwrite site): dense → global, resolved by name against the target project's
config. Metrics stay in dense space; only the write boundary expands.
localizereads thelabel_spacemarker to decide whether thepredictsegmentation already holds globals (copick mode → no value remap) or the model's
dense labels (legacy → remap as before).
Examples
Assume a project with globals
membrane = 1,gas-vesicle-shell = 2,gas-vesicle = 3, and that we train ongas-vesicle-shellandgas-vesicle.1. Default (
--label-space model) — legacy, unchanged.Targets are painted
gas-vesicle-shell → 1,gas-vesicle → 2; the YAML recordslabels: {gas-vesicle-shell: 1, gas-vesicle: 2}with nolabel_spacemarker.predictis written in the same dense space. Byte-for-byte identical to before.2.
--label-space copick— the new canonical case.targetsseg on disk holds globals{0, 2, 3}.labels: {gas-vesicle-shell: 1, gas-vesicle: 2}(dense channels,minted by ascending global label) and
label_space: copick.{0, 2, 3} → {0, 1, 2}viaorig=[0,1,2,3] → target=[0,0,1,2](membraneis painted but unselected →0).{0, 1, 2} → {0, 2, 3}, so the storedpredictseg isordinary copick multilabel and opens correctly under the project config alone.
3. Fine-tuning (
--label-space copick --from-model pretrained/model_config.yaml).Model channels are inherited per object by name from the pretrained config
instead of being minted fresh, so the channel↔object binding is preserved across
the fine-tune. If the pretrained model bound
gas-vesicle → channel 5, the newtargets record channel
5forgas-vesicle. Requesting an object the pretrainedmodel doesn't define is a hard error (fine-tuning onto a new object set is out of
scope).
4. Cross-project inference.
Train on project A (
gas-vesicle = 3), then runoctopi segment -c projectBwhere
gas-vesicle = 7. Because_build_dense_to_globalresolves each channelby name against project B, the written
predictseg uses7. If project B doesnot define a predicted object, inference fails loudly rather than writing a
meaningless label.
Backward compatibility
The default path is untouched: without
--label-space copick, no marker iswritten and every code path behaves exactly as before.
For existing artifacts (targets YAMLs, model configs, and stored
segmentations produced before this change), there is no
label_spacemarker, so:get_class_inforeportslabel_space = None;build_compaction_mapreturns(None, None), soget_transformsskipsMapLabelValuedentirely — training is a no-op remap;_build_dense_to_globalfalls back to identity (np.arange) — predictions arewritten in the model's dense space, unchanged;
localizetakes the legacy global→dense remap branch.The result is bit-identical behavior for all legacy models, targets, and
weights. No data migration is required — existing files keep working via the
identity path, and the new
.get('labels')/.get('label_space')guards alsoremove a latent
KeyErrorthat could occur on label-less configs.Two consistency checks fail loudly rather than silently mislabeling:
MultiCopickDataModule) derives a single sharedmap from the first session, so all configs must assign the same global label to
each shared object; a disagreement raises
ValueError.labelsandlabel_space.